package com.ztom.daily.n2022.m4;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Random;

/**
 * 随机数索引
 * <p>
 * https://leetcode-cn.com/problems/random-pick-index/
 *
 * @author ZhangTao
 */
public class Day25Pick {

    /**
     * Your Solution object will be instantiated and called as such:
     * Solution obj = new Solution(nums);
     * int param_1 = obj.pick(target);
     */
    class Solution1 {

        private Map<Integer, List<Integer>> map;
        private Random random;

        public Solution1(int[] nums) {
            this.map = new HashMap<>();
            this.random = new Random();
            for (int i = 0; i < nums.length; i++) {
                map.computeIfAbsent(nums[i], k -> new ArrayList<>()).add(i);
            }
        }

        public int pick(int target) {
            List<Integer> list = map.get(target);
            return list.get(random.nextInt(list.size()));
        }
    }

    class Solution {

        private int[] nums;
        private Random random;

        public Solution(int[] nums) {
            this.nums = nums;
            this.random = new Random();
        }

        public int pick(int target) {
            int res = 0;
            // 蓄水池算法, 第 i 号球
            int count = 0;
            for (int i = 0; i < nums.length; i++) {
                // 将相同数值的放到同一个盒子
                if (nums[i] == target) {
                    count++;
                    if (random.nextInt(count) == 0) {
                        res = i;
                    }
                }
            }
            return res;
        }
    }
}
